"""
lc200
200. 岛屿数量
给你一个由 '1'（陆地）和 '0'（水）组成的的二维网格，请你计算网格中岛屿的数量。

岛屿总是被水包围，并且每座岛屿只能由水平方向和/或竖直方向上相邻的陆地连接形成。

此外，你可以假设该网格的四条边均被水包围。
"""
from typing import List


class Solution:
    """Solution."""

    @staticmethod
    def numIslands(grid: List[List[int]]) -> int:
        """number of islands."""

        def dfs(grid, i, j, n, m):
            """dfs.

            Args:
                grid:
                i:
                j:
                n:
                m:
            """
            if grid[i][j] == "0":
                return
            grid[i][j] = "0"
            if i > 0 and j > 0:
                dfs(grid, i - 1, j - 1, n, m)
            if i > 0 and j < m - 1:
                dfs(grid, i - 1, j + 1, n, m)
            if i < n - 1 and j < m - 1:
                dfs(grid, i + 1, j + 1, n, m)
            if i < n - 1 and j > 0:
                dfs(grid, i + 1, j - 1, n, m)
            return

        n = len(grid)
        m = len(grid[0])

        num_islands = 0
        for i in range(n):
            for j in range(m):
                if grid[i][j] == "1":
                    num_islands += 1
                    dfs(grid, i, j, n, m)
        return num_islands

    def numIsland_better(grid: List[List[int]]) -> int:
        """numIsland_better.

        Args:
            grid (List[List[str]]): grid

        Returns:
            int:
        """
        nr = len(grid)
        if nr == 0:
            return 0
        nc = len(grid[0])

        uf = UnionFind(grid)
        for r in range(nr):
            for c in range(nc):
                if grid[r][c] == "1":
                    grid[r][c] = "0"
                    for x, y in [(r - 1, c), (r + 1, c), (r, c - 1), (r, c + 1)]:
                        if 0 <= x < nr and 0 <= y < nc and grid[x][y] == "1":
                            uf.union(r * nc + c, x * nc + y)

        return uf.getCount()


class UnionFind:
    """UnionFind."""

    def __init__(self, grid):
        """__init__.

        Args:
            grid:
        """
        m, n = len(grid), len(grid[0])
        self.count = 0
        self.parent = [-1] * (m * n)
        self.rank = [0] * (m * n)
        for i in range(m):
            for j in range(n):
                if grid[i][j] == "1":
                    self.parent[i * n + j] = i * n + j
                    self.count += 1

    def find(self, i):
        """find.

        Args:
            i:
        """

        if self.parent[i] != i:
            self.parent[i] = self.find(self.parent[i])
        return self.parent[i]

    def union(self, x, y):
        """union.

        Args:
            x:
            y:
        """
        rootx = self.find(x)
        rooty = self.find(y)
        if rootx != rooty:
            if self.rank[rootx] < self.rank[rooty]:
                rootx, rooty = rooty, rootx
            self.parent[rooty] = rootx
            if self.rank[rootx] == self.rank[rooty]:
                self.rank[rootx] += 1
            self.count -= 1

    def getCount(self):
        """getCount."""
        return self.count


grid = [
    ["1", "1", "0", "0", "0"],
    ["1", "1", "0", "0", "0"],
    ["0", "0", "1", "0", "0"],
    ["0", "0", "0", "1", "1"],
]

res = Solution.numIsland_better(grid)
print(res)
